import pandas as pd
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from config import Constant
from torch.utils.data import Dataset
import datetime
from dateutil.relativedelta import relativedelta


def generate_samples(raw_data: np.ndarray, target: np.ndarray, raw_tgt_in: np.ndarray, is_test=False, raw_date=None):
    """利用滑窗生成样本"""
    res_data = []
    res_label = []
    res_tgt_in = []
    # res_date = []
    length = raw_data.shape[0]
    seen = Constant.INPUT_WINDOW
    blind = Constant.OUTPUT_WINDOW
    for i in range(length - seen - blind):
        data = raw_data[i:i + seen]
        label = target[i + blind:i + seen + blind]
        tgt_in = raw_tgt_in[i:i + seen]
        # date = raw_date[i + blind:i + seen + blind]
        # res_date.append(date)
        res_data.append(data)
        res_label.append(label)
        res_tgt_in.append(tgt_in)
    if is_test:
        res_data.append(raw_data[length - seen:length])
        res_tgt_in.append(raw_tgt_in[length - seen:length])
        label = target[length - seen + blind:length]
        label = np.append(label, np.array([0]))
        label = label.reshape(-1, 1)
        res_label.append(label)
        # date = raw_date[length - seen + blind:length]
        # date = np.append(date, date[-1] + relativedelta(months=+1))
        # res_date.append(date)
    return np.array(res_data, dtype=float), np.array(res_label, dtype=float), np.array(res_tgt_in, dtype=float)


def scalar(data):
    """归一化处理"""
    mms = MinMaxScaler(feature_range=(-1, 1))
    res = mms.fit_transform(data)
    return res, mms


def generate_dataset():
    """土壤湿度预测数据集生成"""
    humidity_depth = Constant.DEPTH
    humidity_col = ['year', humidity_depth]
    # humidity_col = ['year', '10', '40', '100', '200']
    humidity_raw = pd.read_excel('data/humidity.xls', names=['year', 'lon', 'lat', '10', '40', '100', '200'])
    humidity_raw = humidity_raw[humidity_col]
    steam_col = ['year', 'mon', 'wm2', 'steam_mm']
    steam_raw = pd.read_excel('data/steam.xls', names=['mon', 'year', 'lon', 'lat', 'wm2', 'steam_mm'])
    weather_use_col = ['降水量(mm)', '降水天数']
    weather_col = ['raw_mm', 'rain_day']

    weather = pd.DataFrame(columns=weather_col)
    steam = pd.DataFrame(columns=steam_col)
    humidity = pd.DataFrame(columns=humidity_col)
    years = [2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022]
    for year in years:
        humidity = humidity.append(humidity_raw[humidity_raw.year == year], ignore_index=True)
        steam = steam.append(steam_raw[steam_raw.year == year], ignore_index=True)
        weather_raw = pd.read_excel(f'data/weather/{year}年.xls', usecols=weather_use_col)
        weather_raw.columns = weather_col
        weather = weather.append(weather_raw, ignore_index=True)
    # 数据整合
    for col in steam_col:
        humidity[col] = steam[col]
    for col in weather_col:
        humidity[col] = weather[col]
    # humidity.to_csv('data/merge_data.csv', index=False)
    years = humidity[['year', 'mon']].values
    date = [datetime.date(ym[0], ym[1], 1) for ym in years]
    humidity = humidity.drop(['year'], axis=1)

    data = humidity.to_numpy().copy()
    target = humidity[humidity_depth].to_numpy().copy()
    target = target.reshape(target.shape + (1,))
    humidity[humidity_depth] = humidity[humidity_depth].shift(1, fill_value=0)  # 生成tgt，即后移一位
    tgt_in = humidity.to_numpy()
    # 归一化处理
    data, _ = scalar(data)
    target, mms = scalar(target)
    tgt_in, _ = scalar(tgt_in)
    # 划分训练集和测试
    div = int(data.shape[0] * 0.8)
    train_data = data[:div]
    test_data = data[div:]
    train_tgt = target[:div]
    test_tgt = target[div:]
    train_tgt_in = tgt_in[:div]
    test_tgt_in = tgt_in[div:]

    # 滑窗生成数据
    train_data, train_label, train_tgt_in = generate_samples(train_data, train_tgt, train_tgt_in)
    test_data, test_label, test_tgt_in = generate_samples(test_data, test_tgt, test_tgt_in, is_test=True)
    print(f'train_data shape:{train_data.shape}, train_label shape:{train_label.shape}, tgt:{train_tgt_in.shape}')

    print(f'共有训练集：{train_data.shape[0]}条，测试集{test_data.shape[0]}条')
    return (train_data, train_label, train_tgt_in), (test_data, test_label, test_tgt_in), mms


class MyDataset(Dataset):
    """数据集类"""
    def __init__(self, ds, has_tgt=True):
        self.data = ds[0]
        self.label = ds[1]
        self.has_tgt = has_tgt
        if has_tgt:
            self.tgt_in = ds[2]

    def __getitem__(self, item):
        if self.has_tgt:
            return self.data[item], self.label[item], self.tgt_in[item]
        else:
            return self.data[item], self.label[item]

    def __len__(self):
        return self.data.shape[0]


if __name__ == "__main__":
    generate_dataset()
